#!/usr/bin/env python3
"""
Fix future dates in stock data CSV files.
Adjusts dates so the most recent date is today or a recent past date.
"""

import pandas as pd
from datetime import datetime, timedelta
from pathlib import Path
import argparse


def fix_dates_in_file(file_path, target_end_date=None):
    """
    Fix dates in a single CSV file.
    
    Args:
        file_path: Path to the CSV file
        target_end_date: Desired end date for the data (default: yesterday)
    """
    if target_end_date is None:
        # Default to yesterday (markets are usually closed today)
        target_end_date = datetime.now() - timedelta(days=1)
    
    print(f"Processing {file_path}...")
    
    try:
        # Read the CSV
        df = pd.read_csv(file_path)
        
        # Check if Date column exists
        if 'Date' not in df.columns:
            print(f"  No 'Date' column found in {file_path}")
            return False
            
        # Parse dates
        df['Date'] = pd.to_datetime(df['Date'])
        
        # Get the date range
        max_date = df['Date'].max()
        min_date = df['Date'].min()
        
        print(f"  Original date range: {min_date.date()} to {max_date.date()}")
        
        # Calculate the offset needed
        days_offset = (max_date - target_end_date).days
        
        if days_offset > 0:
            # Need to shift dates backward
            df['Date'] = df['Date'] - timedelta(days=days_offset)
            
            new_max = df['Date'].max()
            new_min = df['Date'].min()
            print(f"  Adjusted date range: {new_min.date()} to {new_max.date()}")
            
            # Save the file
            df['Date'] = df['Date'].dt.strftime('%Y-%m-%d')
            df.to_csv(file_path, index=False)
            print(f"  ✓ Fixed and saved")
            return True
        else:
            print(f"  ✓ Dates are already valid (no future dates)")
            return False
            
    except Exception as e:
        print(f"  ✗ Error processing {file_path}: {e}")
        return False


def fix_all_stock_data(base_dir='data/stocks', target_date=None):
    """
    Fix dates in all stock data files.
    
    Args:
        base_dir: Base directory containing stock data
        target_date: Target end date for all files
    """
    base_path = Path(base_dir)
    if not base_path.exists():
        print(f"Directory {base_dir} does not exist")
        return
    
    # Process all CSV files recursively
    csv_files = list(base_path.rglob('*.csv'))
    
    if not csv_files:
        print(f"No CSV files found in {base_dir}")
        return
    
    print(f"Found {len(csv_files)} CSV files to process\n")
    
    fixed_count = 0
    for csv_file in csv_files:
        if fix_dates_in_file(csv_file, target_date):
            fixed_count += 1
    
    print(f"\n{'='*50}")
    print(f"Summary: Fixed {fixed_count} out of {len(csv_files)} files")


def main():
    parser = argparse.ArgumentParser(description='Fix future dates in stock data CSV files')
    parser.add_argument(
        '--dir', 
        default='data/stocks',
        help='Base directory containing stock data (default: data/stocks)'
    )
    parser.add_argument(
        '--date',
        help='Target end date in YYYY-MM-DD format (default: yesterday)'
    )
    parser.add_argument(
        '--file',
        help='Process a single file instead of all files'
    )
    
    args = parser.parse_args()
    
    # Parse target date if provided
    target_date = None
    if args.date:
        try:
            target_date = datetime.strptime(args.date, '%Y-%m-%d')
        except ValueError:
            print(f"Invalid date format: {args.date}. Use YYYY-MM-DD format.")
            return
    
    if args.file:
        # Process single file
        fix_dates_in_file(args.file, target_date)
    else:
        # Process all files in directory
        fix_all_stock_data(args.dir, target_date)


if __name__ == '__main__':
    main()